# model.py
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn import preprocessing
# 在文件开头的导入部分添加
import matplotlib.pyplot as plt

import datetime
import warnings

# 忽略警告信息
warnings.filterwarnings("ignore")

# --- 1. 定义模型类 ---
class TemperaturePredictor(nn.Module):
    def __init__(self, input_size, hidden_size=128, output_size=1):
        super(TemperaturePredictor, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.Sigmoid(),
            nn.Linear(hidden_size, output_size)
        )

    def forward(self, x):
        return self.network(x)

# --- 2. 数据加载和预处理函数 ---
def load_and_preprocess_data(filepath='data\\temps.csv'):
    """
    加载并预处理数据
    :param filepath: 数据文件路径
    :return: features_scaled (标准化后的特征), labels (目标值), feature_list (特征名列表), scaler (标准化器)
    """
    # 加载数据
    features = pd.read_csv(filepath)
    print(f"数据加载完成，维度: {features.shape}")
    # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    # 添加处理日期数据的代码 (来自原始 Notebook)
    # 分别得到年，月，日
    years = features['year']
    months = features['month']
    days = features['day']

    # datetime格式
    dates = [str(int(year)) + '-' + str(int(month)) + '-' + str(int(day)) for year, month, day in
             zip(years, months, days)]
    dates = [datetime.datetime.strptime(date, '%Y-%m-%d') for date in dates]
    print("日期数据处理完成")
    # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

    # 独热编码
    features = pd.get_dummies(features)
    feature_list = list(features.columns)
    print(f"独热编码后，特征数量: {len(feature_list)}")

    # 分离标签和特征
    labels = np.array(features['actual'])
    features = features.drop('actual', axis=1)
    feature_list.remove('actual')
    features = np.array(features)
    print(f"特征数组形状: {features.shape}, 标签数组形状: {labels.shape}")

    # 特征标准化
    scaler = preprocessing.StandardScaler().fit(features)
    features_scaled = scaler.transform(features)
    print("特征标准化完成")
    # 返回原始特征 DataFrame 用于绘图
    original_features_df = pd.read_csv(filepath)
    return features_scaled, labels, feature_list, scaler,dates,original_features_df

# --- 3. 训练模型函数 ---
def train_model(features, labels, input_size, hidden_size=128, output_size=1,
                batch_size=16, learning_rate=0.001, epochs=1000, print_interval=100):
    """
    训练模型
    :param features: 标准化后的特征
    :param labels: 目标值
    :param input_size: 输入特征维度
    :param hidden_size: 隐藏层大小
    :param output_size: 输出维度
    :param batch_size: 批次大小
    :param learning_rate: 学习率
    :param epochs: 训练轮数
    :param print_interval: 打印间隔
    :return: 训练好的模型
    """
    # 定义模型、损失函数和优化器
    model = TemperaturePredictor(input_size, hidden_size, output_size)
    criterion = nn.MSELoss(reduction='mean')
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    print("开始训练模型...")
    losses = []
    for i in range(epochs):
        epoch_loss = []
        # MINI-Batch训练
        for start in range(0, len(features), batch_size):
            end = min(start + batch_size, len(features))
            # 转换为Tensor
            x_batch = torch.tensor(features[start:end], dtype=torch.float, requires_grad=True)
            y_batch = torch.tensor(labels[start:end], dtype=torch.float, requires_grad=True).view(-1, 1) # 确保标签是二维的

            # 前向传播
            predictions = model(x_batch)
            loss = criterion(predictions, y_batch)

            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            epoch_loss.append(loss.item()) # 使用 .item() 获取数值

        # 记录损失
        if i % print_interval == 0:
            avg_loss = np.mean(epoch_loss)
            losses.append(avg_loss)
            print(f'Epoch [{i}/{epochs}], Loss: {avg_loss:.4f}')

    print("模型训练完成")
    return model

# --- 4. 保存模型函数 ---
def save_model(model, scaler, feature_list, model_path='temperature_model.pth', scaler_path='scaler.pkl', feature_list_path='feature_list.pkl'):
    """
    保存模型、标准化器和特征列表
    :param model: 训练好的模型
    :param scaler: 标准化器
    :param feature_list: 特征名列表
    :param model_path: 模型保存路径
    :param scaler_path: 标准化器保存路径
    :param feature_list_path: 特征列表保存路径
    """
    torch.save(model.state_dict(), model_path)
    import pickle
    with open(scaler_path, 'wb') as f:
        pickle.dump(scaler, f)
    with open(feature_list_path, 'wb') as f:
        pickle.dump(feature_list, f)
    print(f"模型已保存至 {model_path}")
    print(f"标准化器已保存至 {scaler_path}")
    print(f"特征列表已保存至 {feature_list_path}")

# --- 5. 加载模型函数 ---
def load_model(model_path='temperature_model.pth', scaler_path='scaler.pkl', feature_list_path='feature_list.pkl', input_size=None, hidden_size=128, output_size=1):
    """
    加载模型、标准化器和特征列表
    :param model_path: 模型保存路径
    :param scaler_path: 标准化器保存路径
    :param feature_list_path: 特征列表保存路径
    :param input_size: 输入特征维度 (加载模型时需要)
    :param hidden_size: 隐藏层大小
    :param output_size: 输出维度
    :return: model, scaler, feature_list
    """
    import pickle
    with open(scaler_path, 'rb') as f:
        scaler = pickle.load(f)
    with open(feature_list_path, 'rb') as f:
        feature_list = pickle.load(f)

    model = TemperaturePredictor(input_size, hidden_size, output_size)
    model.load_state_dict(torch.load(model_path))
    model.eval() # 设置为评估模式
    print(f"模型已从 {model_path} 加载")
    print(f"标准化器已从 {scaler_path} 加载")
    print(f"特征列表已从 {feature_list_path} 加载")
    return model, scaler, feature_list

# --- 6. 预测函数 ---
def predict_single_sample(model, scaler, feature_list, input_dict):
    """
    对单个样本进行预测
    :param model: 训练好的模型
    :param scaler: 标准化器
    :param feature_list: 特征名列表
    :param input_dict: 包含所有特征的字典，例如:
                       {
                           'year': 2016, 'month': 1, 'day': 1,
                           'temp_2': 0, 'temp_1': 0, 'average': 45,
                           'week_Fri': 0, 'week_Mon': 0, 'week_Sat': 0,
                           'week_Sun': 0, 'week_Thu': 0, 'week_Tue': 0, 'week_Wed': 1,
                           'friend': 45
                       }
    :return: 预测的温度值
    """
    # 1. 构造特征向量
    try:
        input_features = np.array([input_dict[feature] for feature in feature_list])
    except KeyError as e:
        raise KeyError(f"输入字典缺少特征: {e}")

    # 2. 标准化
    input_features_scaled = scaler.transform(input_features.reshape(1, -1)) # reshape 为二维数组

    # 3. 转换为Tensor
    input_tensor = torch.tensor(input_features_scaled, dtype=torch.float)

    # 4. 预测
    with torch.no_grad(): # 预测时不需要计算梯度
        prediction = model(input_tensor)

    # 5. 返回结果
    return prediction.item()


# --- 添加新的绘图函数 ---
def plot_features(dates, features_df):
    """
    绘制特征随时间变化的图表
    :param dates: 日期列表
    :param features_df: 原始特征 DataFrame (未进行独热编码)
    """
    print("开始绘制特征图表...")
    try:
        # 指定默认风格
        plt.style.use('fivethirtyeight')

        # 设置布局
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(nrows=2, ncols=2, figsize=(10, 10))
        fig.autofmt_xdate(rotation=45)

        # 标签值 (实际温度)
        ax1.plot(dates, features_df['actual'])
        ax1.set_xlabel('');
        ax1.set_ylabel('Temperature');
        ax1.set_title('Max Temp')

        # 昨天
        ax2.plot(dates, features_df['temp_1'])
        ax2.set_xlabel('');
        ax2.set_ylabel('Temperature');
        ax2.set_title('Previous Max Temp')

        # 前天
        ax3.plot(dates, features_df['temp_2'])
        ax3.set_xlabel('Date');
        ax3.set_ylabel('Temperature');
        ax3.set_title('Two Days Prior Max Temp')

        # 我的逗逼朋友
        ax4.plot(dates, features_df['friend'])
        ax4.set_xlabel('Date');
        ax4.set_ylabel('Temperature');
        ax4.set_title('Friend Estimate')

        plt.tight_layout(pad=2)

        # 显示图形
        plt.show()
        print("特征图表绘制完成并已显示。")
    except Exception as e:
        print(f"绘图过程中发生错误: {e}")

# --- 7. 主执行逻辑 ---
if __name__ == '__main__':
    # --- 数据和模型参数 ---
    DATA_FILE = 'data\\temps.csv'
    MODEL_PATH = 'temperature_model.pth'
    SCALER_PATH = 'scaler.pkl'
    FEATURE_LIST_PATH = 'feature_list.pkl'
    HIDDEN_SIZE = 128
    OUTPUT_SIZE = 1
    BATCH_SIZE = 16
    LEARNING_RATE = 0.001
    EPOCHS = 1000
    PRINT_INTERVAL = 100

    # --- 加载和预处理数据 ---
    # features_scaled, labels, feature_list, scaler = load_and_preprocess_data(DATA_FILE)
    # input_size = features_scaled.shape[1]
    # --- 加载和预处理数据 (注意返回值增加了 dates 和 original_features_df) ---
    # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    features_scaled, labels, feature_list, scaler, dates, original_features_df = load_and_preprocess_data(DATA_FILE)
    # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
    input_size = features_scaled.shape[1]
    # --- 绘制特征图表 ---
    plot_features(dates, original_features_df)
    # --- 训练模型 ---
    model = train_model(features_scaled, labels, input_size, HIDDEN_SIZE, OUTPUT_SIZE,
                        BATCH_SIZE, LEARNING_RATE, EPOCHS, PRINT_INTERVAL)

    # --- 保存模型 ---
    save_model(model, scaler, feature_list, MODEL_PATH, SCALER_PATH, FEATURE_LIST_PATH)

    # --- 加载模型 (演示用) ---
    # loaded_model, loaded_scaler, loaded_feature_list = load_model(MODEL_PATH, SCALER_PATH, FEATURE_LIST_PATH, input_size, HIDDEN_SIZE, OUTPUT_SIZE)

    # --- 使用模型进行预测 (演示用) ---
    # 示例输入数据 (需要提供所有特征)
    # 注意：week_* 字段是独热编码后的结果，只有一个为1，其余为0
    sample_input = {
        'year': 2016, 'month': 1, 'day': 1,
        'temp_2': 0, 'temp_1': 0, 'average': 45,
        'week_Fri': 0, 'week_Mon': 0, 'week_Sat': 0,
        'week_Sun': 0, 'week_Thu': 0, 'week_Tue': 0, 'week_Wed': 1, # 假设是周三
        'friend': 45
    }
    # 确保字典包含所有特征
    # 这里我们用已有的数据来构造一个完整的输入
    # 从原始 features DataFrame 中取一行作为示例
    original_features_df = pd.read_csv(DATA_FILE)
    original_features_df = pd.get_dummies(original_features_df)
    example_row = original_features_df.iloc[0].drop('actual').to_dict()
    print("\n--- 示例预测 ---")
    print(f"输入特征: {example_row}")

    # 使用训练好的模型预测
    predicted_temp = predict_single_sample(model, scaler, feature_list, example_row)
    actual_temp = original_features_df.iloc[0]['actual'] # 获取真实值
    print(f"真实温度: {actual_temp}")
    print(f"预测温度: {predicted_temp:.2f}")

    # 使用加载的模型预测 (演示用)
    # predicted_temp_loaded = predict_single_sample(loaded_model, loaded_scaler, loaded_feature_list, example_row)
    # print(f"加载模型预测温度: {predicted_temp_loaded:.2f}")

